#include "search/keyword_recommend.hpp"
#include "external/utfcpp/utf8.h"
#include <fstream>
#include <sstream>
#include <algorithm>
#include <climits>
#include <iostream>
#include <stdexcept> // 需要包含这个头文件
#include <tuple>     // 需要包含这个头文件
#include <set>


KeywordRecommend::KeywordRecommend() {}

bool KeywordRecommend::init(const std::string& cnDictPath, const std::string& cnIndexPath,
                          const std::string& enDictPath, const std::string& enIndexPath) {
    if (!loadCNDict(cnDictPath)) return false;
    if (!loadCNIndex(cnIndexPath)) return false;
    if (!loadENDict(enDictPath)) return false;
    if (!loadENIndex(enIndexPath)) return false;
    return true;
}

std::vector<std::pair<std::string, int>> KeywordRecommend::recommend(const std::string& keyword, int count) {
    std::lock_guard<std::mutex> lock(mutex_);
    std::vector<std::pair<std::string, int>> candidates;
    if (keyword.empty()) return candidates;

    std::set<std::string> word_set;
    
    if (isChinese(keyword)) {
        std::vector<std::string> chars = splitChinese(keyword);
        if (chars.empty()) return candidates;
        for(const auto& ch : chars) {
            auto it = cnIndex_.find(ch);
            if(it != cnIndex_.end()){
                word_set.insert(it->second.begin(), it->second.end());
            }
        }
    } else {
        std::string lower_keyword = keyword;
        std::transform(lower_keyword.begin(), lower_keyword.end(), lower_keyword.begin(), ::tolower);
        for(char c : lower_keyword){
            auto it = enIndex_.find(c);
            if(it != enIndex_.end()){
                word_set.insert(it->second.begin(), it->second.end());
            }
        }
    }

    std::vector<std::tuple<int, int, std::string>> sorted_candidates;
    for(const auto& word : word_set) {
        int dist = editDistance(keyword, word);
        int freq = 0;
        if(isChinese(word)){
            auto it = cnDict_.find(word);
            if(it != cnDict_.end()) freq = it->second;
        } else {
            auto it = enDict_.find(word);
            if(it != enDict_.end()) freq = it->second;
        }
        sorted_candidates.emplace_back(dist, freq, word);
    }
    
    std::sort(sorted_candidates.begin(), sorted_candidates.end(), 
        [](const auto& a, const auto& b) {
            if (std::get<0>(a) != std::get<0>(b)) return std::get<0>(a) < std::get<0>(b);
            if (std::get<1>(a) != std::get<1>(b)) return std::get<1>(a) > std::get<1>(b);
            return std::get<2>(a) < std::get<2>(b);
    });
    
    candidates.reserve(count);
    for (const auto& cand : sorted_candidates) {
        if(candidates.size() < (size_t)count) {
             candidates.emplace_back(std::get<2>(cand), std::get<1>(cand));
        } else {
            break;
        }
    }
    
    return candidates;
}

bool KeywordRecommend::isChinese(const std::string& str) const {
    if (str.empty()) return false;
    uint32_t codepoint;
    auto it = str.begin();
    try {
        codepoint = utf8::peek_next(it, str.end());
    } catch(...) { return false; }
    return (codepoint >= 0x4E00 && codepoint <= 0x9FA5);
}

std::vector<std::string> KeywordRecommend::splitChinese(const std::string& str) const {
    std::vector<std::string> chars;
    if (str.empty()) return chars;
    auto it = str.begin();
    auto end = str.end();
    while (it < end) {
        auto start = it;
        try {
            utf8::next(it, end);
            chars.emplace_back(start, it);
        } catch(...) { break; }
    }
    return chars;
}

int KeywordRecommend::editDistance(const std::string& s1, const std::string& s2) const {
    const size_t m = s1.size();
    const size_t n = s2.size();
    if (m == 0) return n;
    if (n == 0) return m;

    std::vector<int> dp(n + 1);
    for (size_t j = 0; j <= n; ++j) dp[j] = j;

    for (size_t i = 1; i <= m; ++i) {
        int prev_val = dp[0]; // 保存 dp[i-1][j-1]
        dp[0] = i;
        for (size_t j = 1; j <= n; ++j) {
            int temp = dp[j];
            if (s1[i - 1] == s2[j - 1]) {
                dp[j] = prev_val;
            } else {
                dp[j] = std::min({prev_val, dp[j], dp[j - 1]}) + 1;
            }
            prev_val = temp;
        }
    }
    return dp[n];
}


// --- 修改点在这里 ---
bool KeywordRecommend::loadCNDict(const std::string& path) {
    std::ifstream file(path);
    if (!file.is_open()) {
        std::cerr << "无法打开中文词典文件: " << path << std::endl;
        return false;
    }
    std::string line, word, freq_str;
    while (std::getline(file, line)) {
        if (line.empty()) continue; // **增加空行检查**
        std::stringstream ss(line);
        if (ss >> word >> freq_str) { // **确保能读出两个部分**
            try {
                cnDict_[word] = std::stoi(freq_str);
            } catch (const std::invalid_argument& e) {
                std::cerr << "警告: 忽略格式错误的行 (cn_dict): " << line << std::endl;
            }
        }
    }
    std::cout << "加载中文词典完成，共 " << cnDict_.size() << " 个词" << std::endl;
    return true;
}

bool KeywordRecommend::loadCNIndex(const std::string& path) {
    std::ifstream file(path);
    if (!file.is_open()) {
        std::cerr << "无法打开中文索引文件: " << path << std::endl;
        return false;
    }
    std::string line;
    while (std::getline(file, line)) {
        if (line.empty()) continue; // **增加空行检查**
        std::stringstream ss(line);
        std::string firstChar, word;
        if (ss >> firstChar) { // **确保能读出首字符**
            while(ss >> word) {
                cnIndex_[firstChar].push_back(word);
            }
        }
    }
    std::cout << "加载中文索引完成，共 " << cnIndex_.size() << " 个首字符索引" << std::endl;
    return true;
}

bool KeywordRecommend::loadENDict(const std::string& path) {
    std::ifstream file(path);
    if (!file.is_open()) {
        std::cerr << "无法打开英文词典文件: " << path << std::endl;
        return false;
    }
    std::string line, word, freq_str;
    while (std::getline(file, line)) {
        if (line.empty()) continue; // **增加空行检查**
        std::stringstream ss(line);
        if (ss >> word >> freq_str) { // **确保能读出两个部分**
            try {
                enDict_[word] = std::stoi(freq_str);
            } catch (const std::invalid_argument& e) {
                 std::cerr << "警告: 忽略格式错误的行 (en_dict): " << line << std::endl;
            }
        }
    }
    std::cout << "加载英文词典完成，共 " << enDict_.size() << " 个词" << std::endl;
    return true;
}

bool KeywordRecommend::loadENIndex(const std::string& path) {
    std::ifstream file(path);
    if (!file.is_open()) {
        std::cerr << "无法打开英文索引文件: " << path << std::endl;
        return false;
    }
    std::string line;
    while (std::getline(file, line)) {
        if (line.empty()) continue; // **增加空行检查**
        std::stringstream ss(line);
        char firstChar;
        std::string word;
        if (ss >> firstChar) { // **确保能读出首字母**
            while(ss >> word) {
                enIndex_[firstChar].push_back(word);
            }
        }
    }
    std::cout << "加载英文索引完成，共 " << enIndex_.size() << " 个首字母索引" << std::endl;
    return true;
}
